package com.ruoyi.framework.encrypt.utils;

import com.alibaba.fastjson2.JSONObject;
import org.apache.commons.codec.binary.Base64;
import org.springframework.util.StringUtils;
import javax.crypto.Cipher;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import java.security.*;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA 加密工具类（非对称加密）
 * 模式：RSA/ECB/OAEPWithSHA-256AndMGF1Padding
 * 密钥长度：2048位（安全且性能均衡）
 * 核心作用：安全传输AES密钥（避免AES密钥明文泄露）
 */
public class RSAUtils {
    // 算法名称
    private static final String ALGORITHM = "RSA";
    // 加密模式+填充方式
    private static final String TRANSFORMATION = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    // 密钥长度（2048位：安全级别足够，单段加密上限214字节，满足AES密钥+IV传输）
    private static final int KEY_SIZE = 2048;
    // PEM格式公钥前缀/后缀
    private static final String PUBLIC_KEY_PREFIX = "-----BEGIN PUBLIC KEY-----";
    private static final String PUBLIC_KEY_SUFFIX = "-----END PUBLIC KEY-----";
    // PEM格式私钥前缀/后缀（服务端内部使用）
    private static final String PRIVATE_KEY_PREFIX = "-----BEGIN PRIVATE KEY-----";
    private static final String PRIVATE_KEY_SUFFIX = "-----END PRIVATE KEY-----";
    // 全局密钥对
    private static KeyPair RSA_KEY_PAIR;
    // OAEP参数规范（与JS端SHA-256哈希+MGF1填充对齐）
    private static final OAEPParameterSpec OAEP_PARAM_SPEC = new OAEPParameterSpec(
            "SHA-256", // 哈希算法
            "MGF1",    // 掩码生成函数
            MGF1ParameterSpec.SHA256, // MGF1哈希算法
            PSource.PSpecified.DEFAULT // 默认PSource参数
    );

    /**
     * 生成RSA密钥对（项目启动时调用）
     * @throws NoSuchAlgorithmException 算法不支持异常
     */
    public static void generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
        keyPairGenerator.initialize(KEY_SIZE); // 初始化密钥长度
        RSA_KEY_PAIR = keyPairGenerator.generateKeyPair();
    }

    /**
     * 对外提供：获取PEM格式公钥
     * @return PEM格式RSA公钥
     */
    public static String getPublicKey() {
        if (RSA_KEY_PAIR == null) {
            throw new RuntimeException("RSA密钥对未初始化");
        }
        return encodePublicKeyToPEM(RSA_KEY_PAIR.getPublic());
    }

    /**
     * 内部使用：获取PEM格式私钥（服务端解密用，不对外暴露）
     * @return PEM格式RSA私钥
     */
    private static String getPrivateKeyPEM() {
        if (RSA_KEY_PAIR == null) {
            throw new RuntimeException("RSA密钥对未初始化");
        }
        return encodePrivateKeyToPEM(RSA_KEY_PAIR.getPrivate());
    }

    /**
     * RSA公钥加密（仅用于加密AES密钥+IV，短数据加密）
     * @param plainText 待加密文本（AES密钥+IV，约32字节）
     * @param publicKey RSA公钥对象
     * @return Base64编码的密文
     * @throws Exception 加密异常
     */
    private static String rsaEncrypt(String plainText, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        // 关键修改：使用OAEP参数规范初始化加密器
        cipher.init(Cipher.ENCRYPT_MODE, publicKey, OAEP_PARAM_SPEC);

        byte[] plainBytes = plainText.getBytes("UTF-8");
        byte[] encryptBytes = cipher.doFinal(plainBytes);
        return Base64.encodeBase64String(encryptBytes);
    }

    /**
     * RSA私钥解密
     * @param cipherTextBase64 Base64编码的RSA密文
     * @param privateKey RSA私钥对象
     * @return 解密后的明文（AES密钥+IV）
     * @throws Exception 解密异常
     */
    private static String rsaDecrypt(String cipherTextBase64, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        // 关键修改：使用OAEP参数规范初始化解密器（与加密端一致）
        cipher.init(Cipher.DECRYPT_MODE, privateKey, OAEP_PARAM_SPEC);

        // 清除密文中的空白字符（避免传输过程中添加的换行/空格）
        String cleanCipherText = cipherTextBase64.replaceAll("\\s+", "");
        byte[] cipherBytes = Base64.decodeBase64(cleanCipherText);

        byte[] decryptBytes = cipher.doFinal(cipherBytes);
        return new String(decryptBytes, "UTF-8");
    }

    /**
     * 从PEM格式公钥解析PublicKey对象
     * @param publicKeyPEM PEM格式公钥
     * @return PublicKey对象
     * @throws Exception 解析异常
     */
    private static PublicKey getPublicKeyFromPEM(String publicKeyPEM) throws Exception {
        if (!StringUtils.hasText(publicKeyPEM)) {
            throw new IllegalArgumentException("公钥不能为空");
        }

        // 去除PEM格式的前缀、后缀和空白字符
        String publicKeyStr = publicKeyPEM.replace(PUBLIC_KEY_PREFIX, "")
                .replace(PUBLIC_KEY_SUFFIX, "")
                .replaceAll("\\s+", "");

        // Base64解码并生成PublicKey对象
        byte[] publicKeyBytes = Base64.decodeBase64(publicKeyStr);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * 从PEM格式私钥解析PrivateKey对象
     * @param privateKeyPEM PEM格式私钥
     * @return PrivateKey对象
     * @throws Exception 解析异常
     */
    private static PrivateKey getPrivateKeyFromPEM(String privateKeyPEM) throws Exception {
        if (!StringUtils.hasText(privateKeyPEM)) {
            throw new IllegalArgumentException("私钥不能为空");
        }

        // 去除PEM格式的前缀、后缀和空白字符
        String privateKeyStr = privateKeyPEM.replace(PRIVATE_KEY_PREFIX, "")
                .replace(PRIVATE_KEY_SUFFIX, "")
                .replaceAll("\\s+", "");

        // Base64解码并生成PrivateKey对象
        byte[] privateKeyBytes = Base64.decodeBase64(privateKeyStr);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * PublicKey对象转PEM格式公钥（带前缀后缀）
     * @param publicKey PublicKey对象
     * @return PEM格式公钥
     */
    private static String encodePublicKeyToPEM(PublicKey publicKey) {
        byte[] publicKeyBytes = publicKey.getEncoded();
        String base64Str = Base64.encodeBase64String(publicKeyBytes);
        // 每64字符换行（标准PEM格式，JS端解析更兼容）
        return PUBLIC_KEY_PREFIX + "\n" + insertLineBreaks(base64Str) + "\n" + PUBLIC_KEY_SUFFIX;
    }

    /**
     * PrivateKey对象转PEM格式私钥（带前缀后缀）
     * @param privateKey PrivateKey对象
     * @return PEM格式私钥
     */
    private static String encodePrivateKeyToPEM(PrivateKey privateKey) {
        byte[] privateKeyBytes = privateKey.getEncoded();
        String base64Str = Base64.encodeBase64String(privateKeyBytes);
        // 每64字符换行（标准PEM格式）
        return PRIVATE_KEY_PREFIX + "\n" + insertLineBreaks(base64Str) + "\n" + PRIVATE_KEY_SUFFIX;
    }

    /**
     * 给Base64字符串添加换行（符合PEM格式规范）
     * @param str Base64字符串
     * @return 带换行的Base64字符串
     */
    private static String insertLineBreaks(String str) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < str.length(); i += 64) {
            int end = Math.min(i + 64, str.length());
            sb.append(str.substring(i, end)).append("\n");
        }
        return sb.toString().trim(); // 去除最后一个换行
    }

    /**
     * 混合加密（服务端加密超长文本）
     * @param plainText 待加密明文（任意长度）
     * @return Map{ encryptedKey: RSA加密后的AES密钥IV, encryptedText: AES加密后的文本 }
     * @throws Exception 加密异常
     */
    public static Map<String, String> encrypt(String plainText) throws Exception {
        if (!StringUtils.hasText(plainText)) {
            throw new IllegalArgumentException("明文不能为空");
        }

        // 1. 生成AES密钥和IV
        Map<String, String> aesInfo = AESUtils.generateAesKeyAndIv();
        String aesKey = aesInfo.get("AES_KEY");
        String aesIv = aesInfo.get("AES_IV");

        // 2. AES加密超长文本（高效处理大量数据）
        String aesCipherText = AESUtils.aesEncrypt(plainText, aesKey, aesIv);

        // 3. 拼接AES密钥和IV（用|分隔，便于解密时拆分）
        String aesKeyIv = aesKey + "|" + aesIv;

        // 4. RSA公钥加密AES密钥+IV（安全传输密钥，使用OAEP填充）
        String rsaCipherKey = rsaEncrypt(aesKeyIv, RSA_KEY_PAIR.getPublic());

        // 5. 返回加密结果
        Map<String, String> result = new HashMap<>();
        result.put("encryptedKey", rsaCipherKey);
        result.put("encryptedText", aesCipherText);
        return result;
    }

    /**
     * 混合解密
     * @param encryptedKey RSA加密后的AES密钥+IV（JS端原生RSA-OAEP加密）
     * @param encryptedText AES加密后的文本密文
     * @return 最终解密后的明文
     * @throws Exception 解密异常
     */
    public static String decrypt(String encryptedKey, String encryptedText) throws Exception {
        // 1. 参数校验
        if (!StringUtils.hasText(encryptedKey) || !StringUtils.hasText(encryptedText)) {
            throw new IllegalArgumentException("encryptedKey和encryptedText不能为空");
        }

        try {
            // 2. RSA私钥解密：获取AES密钥+IV（使用OAEP填充解密）
            String aesKeyIv = rsaDecrypt(encryptedKey, RSA_KEY_PAIR.getPrivate());
            String[] aesInfo = aesKeyIv.split("\\|", 2); // 按|拆分（最多拆2段，避免IV含|）
            if (aesInfo.length != 2) {
                throw new Exception("AES密钥+IV格式错误，无法拆分");
            }
            String aesKey = aesInfo[0];
            String aesIv = aesInfo[1];

            // 3. AES解密：获取原始明文（强制去除PKCS5填充，避免残留符号）
            return AESUtils.aesDecrypt(encryptedText, aesKey, aesIv);
        } catch (Exception e) {
            throw new Exception("混合解密失败：" + e.getMessage(), e);
        }
    }

    public static String decrypt(String encryptedInfo) throws Exception {
        JSONObject jsonObject = JSONObject.parseObject(encryptedInfo);
        String encryptedText = jsonObject.getString("encryptedText");
        String encryptedKey = jsonObject.getString("encryptedKey");
        return decrypt(encryptedKey, encryptedText);
    }

    // ------------------------------ 测试方法 ------------------------------
    public static void main(String[] args) throws Exception {
        generateKeyPair();
        // 测试文本（模拟JS端传输的JSON字符串）
        String testText = "{\"username\":\"admin\",\"password\":\"admin123\",\"code\":\"3\",\"uuid\":\"4d5e40185ec9443891c4b6b4ab00f091\"}";
        System.out.println("测试文本：" + testText);
        System.out.println("测试文本长度：" + testText.length());

        // 1. 服务端混合加密（模拟JS端加密逻辑）
        Map<String, String> encryptResult = encrypt(testText);
        System.out.println("加密后 encryptedKey：" + encryptResult.get("encryptedKey"));
        System.out.println("加密后 encryptedText：" + encryptResult.get("encryptedText"));

        // 2. 服务端混合解密（模拟接收JS端密文后解密）
        String decryptText = decrypt(encryptResult.get("encryptedKey"), encryptResult.get("encryptedText"));
        System.out.println("解密后文本：" + decryptText);
        System.out.println("解密后文本长度：" + decryptText.length());
        System.out.println("加解密一致性：" + testText.equals(decryptText)); // 应输出true
        System.out.println("JSON解析测试：" + JSONObject.parseObject(decryptText)); // 验证无残留符号，可正常解析
    }
}